DEFAULT_INT_NAMES = [f"int_{i}" for i in range(13)]
DEFAULT_CAT_NAMES = [f"cat_{i}" for i in range(26)]

def plan_0():
    with open("preprocess_ir.txt", "w") as f:
        f.write("fillnull int_0 int_0 value=0\n")
        f.write("bucketize int_0 bucketize_int_0 boundaries=0.5,1.0,1.5\n")
        f.write("logit int_0 int_0 eps=1e-5\n")
        f.write("sigrid_hash cat_0 cat_0 salt=salt_value num_buckets=65536\n")
        f.write("ngram cat_0 cat_0 ngram_param=4\n")
        f.write("firstx cat_0 cat_0 count=1\n")
        f.write("ngram bucketize_int_0 bucketize_int_0 ngram_param=4\n")
        f.write("firstx bucketize_int_0 bucketize_int_0 count=1\n")

def plan_1():
    with open("preprocess_ir.txt", "w") as f:

        # Step 1: fillnull + cast
        for name in DEFAULT_INT_NAMES:
            f.write(f"fillnull {name} {name} value=0\n")
            f.write(f"cast {name} {name} dtype=float32\n")

        for name in DEFAULT_CAT_NAMES:
            f.write(f"fillnull {name} {name} value=0\n")
            f.write(f"cast {name} {name} dtype=int64\n")

        # Step 2: bucketize int_0 -> bucketize_int_0
        f.write("bucketize int_0 bucketize_int_0 boundaries=0.5,1.0,1.5\n")
        f.write("cast bucketize_int_0 bucketize_int_0 dtype=int64\n")

        # Step 3: logit for selected ints (assume self.args.int_name = DEFAULT_INT_NAMES)
        for name in DEFAULT_INT_NAMES:
            f.write(f"logit {name} {name} eps=1e-5\n")

        # Step 5: sigrid_hash + array + firstx for selected cats (assume self.args.cat_name = DEFAULT_CAT_NAMES)
        for idx, name in enumerate(DEFAULT_CAT_NAMES):
            f.write(f"sigrid_hash {name} {name} salt=salt_value num_buckets=65536\n")
            f.write(f"firstx {name} {name} count=1\n")

def plan_2():
    
    nDense = 52
    nSparse = 104
    
    with open("preprocess_ir.txt", "w") as f:

        # Step 1: fillnull + cast
        for name in DEFAULT_INT_NAMES:
            f.write(f"fillnull {name} {name} value=0\n")
            f.write(f"cast {name} {name} dtype=float32\n")

        for name in DEFAULT_CAT_NAMES:
            f.write(f"fillnull {name} {name} value=0\n")
            f.write(f"cast {name} {name} dtype=int64\n")

        # Step 2: copy dense/sparse features
        for i in range(nDense):
            src = f"int_{i % 13}"
            dst = f"int_{i}"
            f.write(f"copy {src} {dst}\n")

        for i in range(nSparse):
            src = f"cat_{i % 26}"
            dst = f"cat_{i}"
            f.write(f"copy {src} {dst}\n")

        # Step 3: dense transformations
        for _ in range(4):  # repeat 4 times
            for i in range(0, 16):
                f.write(f"logit int_{i} int_{i} eps=1e-5\n")
            for i in range(16, 37):
                f.write(f"boxcox int_{i} int_{i} eps=1e-5\n")
            for i in range(37, 52):
                f.write(f"onehot int_{i} int_{i} min=0.0 max=16.0 dim=5\n")
                f.write(f"firstx int_{i} int_{i} count=1\n")

        # Step 4: sparse transformations
        for i in range(0, 25):
            f.write(f"sigrid_hash cat_{i} cat_{i} min=0 max=65536\n")

        for i in range(25, 53):
            f.write(f"clamp cat_{i} cat_{i} min=0 max=16\n")

        for i in range(53, 78):
            f.write(f"firstx cat_{i} cat_{i} count=1\n")

        for i in range(78, 104):
            f.write(f"onehot cat_{i} cat_{i} min=0.0 max=16.0 dim=8\n")
            f.write(f"ngram cat_{i} cat_{i} n=4\n")
            f.write(f"firstx cat_{i} cat_{i} count=1\n")

if __name__ == "__main__":
    import sys

    if len(sys.argv) != 2:
        print("Usage: python script.py <plan_number>")
        sys.exit(1)

    plan_number = sys.argv[1]

    if plan_number == "0":
        plan_0()
    elif plan_number == "1":
        plan_1()
    elif plan_number == "2":
        plan_2()
    else:
        print("Invalid plan number. Please choose 0, 1, or 2.")
        sys.exit(1)